package edu.zju.gis.dbfg.model.tile;

import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.common.util.FileUtil;
import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.base.Pipeline;
import edu.zju.gis.dbfg.model.util.GeoItem;
import edu.zju.gis.dbfg.model.util.TileID;
import edu.zju.gis.dbfg.model.util.TileUtil;
import edu.zju.gis.dbfg.model.util.ZLevelInfo;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.PairFlatMapFunction;
import org.apache.spark.broadcast.Broadcast;
import org.geotools.geometry.jts.JTS;
import org.geotools.referencing.CRS;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.geom.GeometryFactory;
import org.locationtech.jts.io.WKTReader;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.File;
import java.util.*;
import java.util.List;
import java.util.stream.Collectors;

import static edu.zju.gis.dbfg.common.constant.TileConstant.SCREEN_TILE_BUFFER;
import static edu.zju.gis.dbfg.common.constant.TileConstant.SCREEN_TILE_SIZE;


/**
 * @author Hu
 * @date 2019/9/4
 * 基于Elasticsearch数据源，实现矢量瓦片构建
 * 目前只支持索引全量数据构建瓦片，不支持数据过滤及字段选择
 * TODO 默认输入数据坐标系为 4326，假设 4326 与 4490 等价
 **/
public class VectorTileClipper {

    private final static Logger logger = LoggerFactory.getLogger(VectorTileClipper.class);

    public static void main(String[] args) throws Exception {

        // Setup environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf();
        // conf.setMaster("local[8]");
        // conf.setAppName("lalala");
        JavaSparkContext jsc = new JavaSparkContext(conf);

        // Setup model arguments
        logger.info("Setup arguments");
        VectorTileClipperArgs mArgs = VectorTileClipperArgs.initArgs(args);
        if (mArgs == null) {
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(VectorTileClipper.class, "main()", msg, args);
        }

        File outDir = new File(mArgs.getOutDir());
        if (!outDir.exists()) {
            outDir.mkdirs();
            logger.info("create new directory for output tiles: " + outDir.getAbsolutePath());
        }

        CoordinateReferenceSystem targetCrs = CRS.parseWKT(mArgs.getTargetCRS());
        CoordinateReferenceSystem sourceCrs = CRS.parseWKT(mArgs.getSourceCRS());

        List<Double> targetCrsExtents = Arrays.stream(mArgs.getExtent().split(",")).map(new java.util.function.Function<String, Double>() {
            @Override
            public Double apply(String s) {
                return Double.valueOf(s);
            }
        }).collect(Collectors.toList());

        // 定义金字塔网格参数
        PyramidConfig pyramidConfig = new PyramidConfig.PyramidConfigBuilder()
                .setBaseMapEnv(targetCrsExtents.get(0), targetCrsExtents.get(1), targetCrsExtents.get(2), targetCrsExtents.get(3))
                .setCrs(targetCrs)
                .setzLevelRange(mArgs.getZMin(), mArgs.getZMax())
                .build();

        WKTReader r = new WKTReader();
        String qgWkt = FileUtil.readByLine(mArgs.getQgWktFile(), 1, false).get(0);
        Geometry qg = r.read(qgWkt);
        GeometryFactory f = new GeometryFactory();

        // Setup data from datasource
        logger.info("Setup datasource");
        String[] datasources = mArgs.getInput().split(",");
        for (int i = 0; i < datasources.length; i++) {
            String datasource = datasources[i];
            if (!datasource.startsWith("es://")) {
                String msg = String.format("Unvalid datasource %s, only support elasticsearch datastore yet (url schema as 'es://{indexname}')", datasource);
                logger.error(msg);
                throw new ModelFailedException(VectorTileClipper.class, "main()", msg, args);
            }
            datasources[i] = datasource.replace("es://", "");
        }

        if (!mArgs.getTargetCRS().equals(mArgs.getSourceCRS())) {
            MathTransform transform = CRS.findMathTransform(sourceCrs, targetCrs);
            qg = JTS.transform(qg, transform);
        }

        Envelope jobEnvelope = qg.getEnvelopeInternal();

        ZLevelInfo[] jobLevels = TileUtil.initZLevelInfoPZ(pyramidConfig, jobEnvelope);

        List<TileID> jobTiles = new ArrayList<>();
        ZLevelInfo jobLevel = jobLevels[mArgs.getJobSplitZLevel() - mArgs.getZMin()];
        int tx_min = jobLevel.getTileRanges()[0];
        int tx_max = jobLevel.getTileRanges()[1];
        int ty_min = jobLevel.getTileRanges()[2];
        int ty_max = jobLevel.getTileRanges()[3];
        for (int x = tx_min; x <= tx_max; x++) {
            for (int y = ty_min; y <= ty_max; y++) {
                TileID t = new TileID();
                t.setX(x);
                t.setY(y);
                t.setzLevel(mArgs.getJobSplitZLevel());
                Envelope en = TileUtil.createTileBox(t, pyramidConfig);
                Geometry eng = f.toGeometry(en);
                boolean flag = eng.intersects(qg);
                if (flag) {
                    jobTiles.add(t);
                }
            }
        }

        logger.info(String.format(" ============ TOTAL TILE NUM: %d =========== ", jobTiles.size()));

        List<Integer> tileIndex = new ArrayList<>();
        for (int i=0; i<jobTiles.size(); i++) {
            tileIndex.add(i);
        }

        TileID[] jobTilesArr = new TileID[jobTiles.size()];

        final Broadcast<VectorTileClipperArgs> mArgsBroad = jsc.broadcast(mArgs);
        final Broadcast<PyramidConfig> pyramidConfigBroad = jsc.broadcast(pyramidConfig);
        final Broadcast<TileID[]> jobTilesBroad = jsc.broadcast(jobTiles.toArray(jobTilesArr));

        JavaPairRDD<TileID, GeoItem> allFeatureRDD = jsc.parallelize(tileIndex, tileIndex.size()).flatMapToPair(new PairFlatMapFunction<Integer, TileID, GeoItem>() {
            @Override
            public Iterator<Tuple2<TileID, GeoItem>> call(Integer index) throws Exception {
                TileID tileID = jobTilesBroad.getValue()[index];
                TileJob tileJob = new TileJob();
                tileJob.setFileType(mArgs.getFileType());
                tileJob.setLayerName(mArgs.getLayerName());
                tileJob.setZMax(mArgs.getZMax());
                tileJob.setOutdir(mArgs.getOutDir());
                tileJob.setPyramidConfig(pyramidConfig);
                tileJob.setSourceCrs(sourceCrs);
                tileJob.setTargetCrs(targetCrs);
                tileJob.setSource(datasources);
                return tileJob.tileJobResult(tileID).iterator();
            }
        });

        /**
         * 迭代计算瓦片各个层级的数据
         */
        for (int i = mArgs.getJobSplitZLevel() - 1; i >= mArgs.getZMin(); i--) {
            allFeatureRDD = allFeatureRDD.groupByKey().flatMapToPair(new PairFlatMapFunction<Tuple2<TileID, Iterable<GeoItem>>, TileID, GeoItem>() {
                @Override
                public Iterator<Tuple2<TileID, GeoItem>> call(Tuple2<TileID, Iterable<GeoItem>> in) throws Exception {

                    // List<GeoItem> features = IteratorUtils.toList(in._2.iterator());

                    Iterator<GeoItem> geoItemIterator = in._2.iterator();

                    TileID tileID = in._1;
                    Envelope tileEnvelope = TileUtil.createTileBox(tileID, pyramidConfig);
                    TileJob tileJob = new TileJob();
                    tileJob.setFileType(mArgs.getFileType());
                    tileJob.setLayerName(mArgs.getLayerName());
                    tileJob.setZMax(mArgs.getZMax());
                    tileJob.setOutdir(mArgs.getOutDir());
                    tileJob.setPyramidConfig(pyramidConfig);
                    tileJob.setSourceCrs(sourceCrs);
                    tileJob.setTargetCrs(targetCrs);
                    tileJob.setSource(datasources);

                    /**
                     * 将之前被分割的图形union回来
                     */
                    Map<String, GeoItem> unionedGeoItem = new HashMap<>();
                    while(geoItemIterator.hasNext()) {
                        GeoItem geoItem = geoItemIterator.next();
                        if (!unionedGeoItem.containsKey(geoItem.getFid())) {
                            unionedGeoItem.put(geoItem.getFid(), geoItem);
                        } else {
                            Geometry g = tileJob.onlyPolygon(geoItem.getGeometry());
                            Geometry og = tileJob.onlyPolygon(unionedGeoItem.get(geoItem.getFid()).getGeometry());
                            if (og == null) continue;
                            geoItem.setGeometry(tileJob.union(g, og, Math.max(tileEnvelope.getWidth()/(16*SCREEN_TILE_SIZE), tileEnvelope.getHeight()/(16*SCREEN_TILE_SIZE)), 0, 2));
                            unionedGeoItem.put(geoItem.getFid(), geoItem);
                        }
                    }

                    /**
                     * 输出当前层级瓦片
                     */
                    tileJob.buildTile(new ArrayList<>(unionedGeoItem.values()), tileID, pyramidConfigBroad.getValue(), mArgsBroad.getValue().getLayerName(), mArgsBroad.getValue().getFileType(), mArgsBroad.getValue().getOutDir());

                    /**
                     * 处理数据用于上一层计算，这时就不需要clip了，只需要进行抽稀即可
                     */
                    TileID upperTileID = new TileID();
                    upperTileID.setzLevel(tileID.getzLevel()-1);
                    upperTileID.setX(tileID.getX()/2);
                    upperTileID.setY(tileID.getY()/2);

                    List<Tuple2<TileID, GeoItem>> result = new ArrayList<>();

                    if (tileID.getzLevel() > mArgs.getZMin()) {
                        Envelope upperLevelEnvelope = TileUtil.createTileBox(upperTileID, pyramidConfig);
                        Pipeline upperLevelSimplifyPipeline = tileJob.getPipeline(targetCrs, SCREEN_TILE_BUFFER, upperLevelEnvelope, false, false, true);
                        Iterator<String> keys = unionedGeoItem.keySet().iterator();
                        while (keys.hasNext()) {
                            String key = keys.next();
                            GeoItem gi = unionedGeoItem.get(key);
                            Geometry g = upperLevelSimplifyPipeline.execute(gi.getGeometry());
                            if (g.isEmpty()) {
                                continue;
                            }
                            g = tileJob.onlyPolygon(g);
                            if (g == null) {
                                continue;
                            }
                            gi.setGeometry(g);
                            result.add(new Tuple2<>(upperTileID, gi));
                        }
                    } else {
                        result.add(new Tuple2<>(null, null));
                    }
                    return result.iterator();
                }
            });

        }

        logger.info("Model Finish: " + allFeatureRDD.count());
        jsc.close();
    }



}
